#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
通义千问多模态API调用工具 (qwen2-vl-72b-gptq-int4)

这个脚本提供了简单实用的函数来调用通义千问的多模态大模型，
支持文本+图像的输入。基于测试结果，该脚本使用图像格式的请求，
即使是纯文本查询也以多模态格式发送，以确保API调用成功。
"""

import os
import base64
import litellm
import requests
import json
from typing import List, Dict, Any, Optional, Union
from rich.console import Console
from rich.markdown import Markdown

# 导入配置和环境变量加载模块
from config_env_load import load_env_variables, load_config, get_model_config, get_api_key

# 初始化控制台
console = Console()

# 加载环境变量
load_env_variables()

# 初始化配置
try:
    config = load_config()
    model_config = get_model_config(config)
    API_BASE = model_config.get("base_url")
    MODEL_NAME = model_config.get("model_name")
    
    # 确保 MODEL_NAME 包含 openai/ 前缀（如果需要）
    if not MODEL_NAME.startswith("openai/") and not MODEL_NAME.startswith("Qwen/"):
        MODEL_NAME = f"openai/{MODEL_NAME}"
        
except Exception as e:
    console.print(f"[bold red]配置加载失败: {str(e)}，使用默认配置[/bold red]")
    # 默认配置，作为备用
    API_BASE = "http://ip:port/v1"
    MODEL_NAME = "openai/qwen2-vl-72b-gptq-int4"  # 使用openai/前缀

def encode_image_to_base64(image_path: str) -> str:
    """
    将图像文件编码为base64字符串
    
    Args:
        image_path: 图像文件路径
        
    Returns:
        base64编码的图像字符串
    """
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def setup_api(verbose: bool = False):
    """
    设置API环境
    
    Args:
        verbose: 是否启用详细日志
    """
    os.environ["OPENAI_API_BASE"] = API_BASE
    
    # 获取API密钥
    api_key = get_api_key(API_BASE)
    os.environ["OPENAI_API_KEY"] = api_key if api_key != "dummy-key" else "fake-api-key"
    litellm.set_verbose = verbose

def query_model(text: str, image_path: Optional[str] = None, stream: bool = False, preserve_numbers: bool = False) -> str:
    """
    查询通义千问模型
    
    Args:
        text: 文本查询
        image_path: 可选的图像路径
        stream: 是否使用流式输出
        preserve_numbers: 是否保持数字原始格式
        
    Returns:
        模型回复的文本
    """
    setup_api()
    
    # 如果需要保持数字原始格式，添加特殊指令
    if preserve_numbers and "提取" in text:
        text = f"{text}\n特别说明: 所有数字必须严格保持原始格式，不进行任何四舍五入、取整或标准化小数位数。"
    
    # 准备消息内容（始终使用多模态格式）
    content = [{"type": "text", "text": text}]
    
    # 如果提供了图像，添加图像内容
    if image_path:
        if image_path.startswith(('http://', 'https://')):
            content.append({"type": "image_url", "image_url": image_path})
        else:
            try:
                base64_image = encode_image_to_base64(image_path)
                content.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{base64_image}"})
            except Exception as e:
                console.print(f"[bold red]图像加载失败: {str(e)}[/bold red]")
                return f"错误: 无法加载图像 {image_path}。错误信息: {str(e)}"
    
    # 准备完整消息
    messages = [{"role": "user", "content": content}]
    
    try:
        if stream:
            # 流式输出
            collected_chunks = []
            for chunk in litellm.completion(
                model=MODEL_NAME,
                messages=messages,
                max_tokens=1024,
                stream=True,
            ):
                content = chunk.choices[0].delta.content
                if content:
                    console.print(content, end="")
                    collected_chunks.append(content)
            
            print()  # 新行
            return "".join(collected_chunks)
        else:
            # 非流式输出
            response = litellm.completion(
                model=MODEL_NAME,
                messages=messages,
                max_tokens=1024,
            )
            return response.choices[0].message.content
    except Exception as e:
        # console.print(f"[bold red]API调用失败: {str(e)}[/bold red]")
        
        # 尝试使用直接API调用作为后备方案
        try:
            console.print("[yellow]尝试使用直接API调用...[/yellow]")
            return direct_api_call(text, image_path, preserve_numbers)
        except Exception as e2:
            console.print(f"[bold red]直接API调用也失败: {str(e2)}[/bold red]")
            return f"错误: 无法获取回复。错误信息: {str(e)}"

def direct_api_call(text: str, image_path: Optional[str] = None, preserve_numbers: bool = False) -> str:
    """
    直接调用API而不使用litellm
    
    Args:
        text: 文本查询
        image_path: 可选的图像路径
        preserve_numbers: 是否保持数字原始格式
        
    Returns:
        模型回复的文本
    """
    url = f"{API_BASE}/chat/completions"
    
    # 获取API密钥
    api_key = get_api_key(API_BASE)
    headers = {"Content-Type": "application/json"}
    if api_key != "dummy-key":
        headers["Authorization"] = f"Bearer {api_key}"
    
    # 如果需要保持数字原始格式，添加特殊指令
    if preserve_numbers and "提取" in text:
        text = f"{text}\n特别说明: 所有数字必须严格保持原始格式，不进行任何四舍五入、取整或标准化小数位数。"
    
    # 准备消息内容（始终使用多模态格式）
    content = [{"type": "text", "text": text}]
    
    # 如果提供了图像，添加图像内容
    if image_path:
        if image_path.startswith(('http://', 'https://')):
            content.append({"type": "image_url", "image_url": image_path})
        else:
            base64_image = encode_image_to_base64(image_path)
            content.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{base64_image}"})
    
    # 从MODEL_NAME中获取真实模型名称（去除可能的'openai/'前缀）
    model_name = MODEL_NAME.split('/')[-1] if '/' in MODEL_NAME else MODEL_NAME
    
    data = {
        "model": model_name,  # 使用配置中的模型名称，但去除前缀
        "messages": [{"role": "user", "content": content}],
        "max_tokens": 1024
    }
    
    response = requests.post(url, headers=headers, data=json.dumps(data))
    if response.status_code == 200:
        result = response.json()
        return result["choices"][0]["message"]["content"]
    else:
        raise Exception(f"API错误: 状态码 {response.status_code}, 回复: {response.text}")

def interactive_chat():
    """运行交互式聊天会话"""
    console.print("[bold green]欢迎使用通义千问多模态聊天！[/bold green]")
    console.print("[bold blue]当前配置:[/bold blue]")
    console.print(f"- 模型: {MODEL_NAME}")
    console.print(f"- API基础URL: {API_BASE}")
    console.print("[bold]命令说明:[/bold]")
    console.print("- 输入 'exit' 或 'quit' 退出")
    console.print("- 输入 'image: 路径' 发送图像")
    console.print("- 输入 'preserve: on' 开启保持数字原始格式模式")
    console.print("- 输入 'preserve: off' 关闭保持数字原始格式模式")
    console.print("- 输入 'reload' 重新加载配置")
    
    preserve_numbers = False
    
    while True:
        try:
            user_input = input("\n[You]: ")
            
            if user_input.lower() in ['exit', 'quit']:
                console.print("[bold]再见！[/bold]")
                break
            
            # 检查是否需要重新加载配置
            if user_input.lower() == 'reload':
                reload_config()
                console.print("[bold green]已重新加载配置[/bold green]")
                console.print(f"- 模型: {MODEL_NAME}")
                console.print(f"- API基础URL: {API_BASE}")
                continue
            
            # 检查是否切换数字保持模式
            if user_input.lower().startswith('preserve:'):
                mode = user_input.lower().split(':', 1)[1].strip()
                if mode == 'on':
                    preserve_numbers = True
                    console.print("[bold green]已开启保持数字原始格式模式[/bold green]")
                elif mode == 'off':
                    preserve_numbers = False
                    console.print("[bold yellow]已关闭保持数字原始格式模式[/bold yellow]")
                continue
            
            # 检查是否发送图像
            image_path = None
            if user_input.startswith('image:'):
                parts = user_input.split(' ', 1)
                if len(parts) > 1:
                    image_path = parts[1].strip()
                    user_input = input("[Image] 请输入您的问题: ")
            
            print("\n[Qwen]: ", end="")
            response = query_model(user_input, image_path, stream=True, preserve_numbers=preserve_numbers)
            
        except KeyboardInterrupt:
            console.print("\n[bold]程序被中断，再见！[/bold]")
            break
        except Exception as e:
            console.print(f"[bold red]发生错误: {str(e)}[/bold red]")

def process_single_query(query: str, image_path: Optional[str] = None):
    """处理单个查询并显示结果"""
    if image_path:
        console.print(f"[bold]查询: {query} [图像: {image_path}][/bold]")
    else:
        console.print(f"[bold]查询: {query}[/bold]")
    
    console.print("[bold blue]正在获取回复...[/bold blue]")
    
    # 检查查询是否与数字提取相关，如果是则启用保持数字格式选项
    preserve_nums = any(keyword in query for keyword in ["提取", "数字", "保持原始格式", "json", "JSON"])
    
    response = query_model(query, image_path, preserve_numbers=preserve_nums)
    
    console.print("\n[bold green]回复:[/bold green]")
    console.print(Markdown(response))

def reload_config():
    """重新加载配置"""
    global API_BASE, MODEL_NAME
    try:
        config = load_config()
        model_config = get_model_config(config)
        API_BASE = model_config.get("base_url")
        MODEL_NAME = model_config.get("model_name")
        
        # 确保 MODEL_NAME 包含 openai/ 前缀（如果需要）
        if not MODEL_NAME.startswith("openai/") and not MODEL_NAME.startswith("Qwen/"):
            MODEL_NAME = f"openai/{MODEL_NAME}"
            
        console.print("[bold green]配置已重新加载[/bold green]")
    except Exception as e:
        console.print(f"[bold red]配置重新加载失败: {str(e)}[/bold red]")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='通义千问多模态API调用工具')
    parser.add_argument('--query', type=str, help='要发送的查询文本')
    parser.add_argument('--image', type=str, help='要发送的图像路径')
    parser.add_argument('--interactive', action='store_true', help='启动交互式聊天模式')
    parser.add_argument('--preserve-numbers', action='store_true', help='保持数字原始格式')
    parser.add_argument('--output', '-o', type=str, help='输出JSON文件路径')
    parser.add_argument('--config', type=str, default='config.yaml', help='配置文件路径')
    
    args = parser.parse_args()
    
    # 如果指定了不同的配置文件，重新加载配置
    if args.config != 'config.yaml':
        try:
            config = load_config(args.config)
            model_config = get_model_config(config)
            API_BASE = model_config.get("base_url")
            MODEL_NAME = model_config.get("model_name")
            
            # 确保 MODEL_NAME 包含 openai/ 前缀（如果需要）
            if not MODEL_NAME.startswith("openai/") and not MODEL_NAME.startswith("Qwen/"):
                MODEL_NAME = f"openai/{MODEL_NAME}"
        except Exception as e:
            console.print(f"[bold red]无法从指定配置文件加载: {str(e)}[/bold red]")
    
    if args.interactive or (not args.query and not args.image):
        interactive_chat()
    else:
        query = args.query or "请描述这张图片，提取的所有数字请保持原始格式，不要统一小数位数"
        # 修改 process_single_query 函数调用，使用 preserve_numbers 参数
        if args.preserve_numbers:
            # 如果查询中没有相关说明，添加保持数字格式的说明
            if not any(keyword in query for keyword in ["原始格式", "数字格式", "小数位"]):
                query = f"{query}。请保持所有数字的原始格式，不进行四舍五入、取整或标准化小数位数。"
        
        process_single_query(query, args.image) 